{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import argparse\n",
    "import copy\n",
    "import json\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "import re\n",
    "import sys\n",
    "from collections import defaultdict, deque\n",
    "from typing import List\n",
    "import re\n",
    "from typing import Any, List\n",
    "import ipydagred3\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import openai\n",
    "from tqdm import tqdm\n",
    "from CHAMP_utils import *\n",
    "\n",
    "#os.environ[\"http_proxy\"] = \"http://localhost:7890\"\n",
    "#os.environ[\"https_proxy\"] = \"http://localhost:7890\"\n",
    "setOpenAi(keyid = 4)\n",
    "\n",
    "GPT_MODEL = \"gpt-4-turbo\"\n",
    "GPT_MODEL = \"gpt-4o-mini\"\n",
    "GPT_MODEL = \"gpt-3.5-turbo\"\n",
    "\n",
    "# 读取问题数据文件\n",
    "file_path = '/Users/natehu/Desktop/Tsinghua Research/ToT复现/TOT/champ_116.json'\n",
    "with open(file_path, 'r', encoding='utf-8') as file:\n",
    "    dataset = json.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'ids': 0,\n",
       "  'problem_identifier': 'P_Combinatorics_3',\n",
       "  'problem_text': 'On a chess board, two rooks are placed peacefully if they are not on the same row or column. For an n x n chess board, find the number of ways that n rooks can be placed peacefully (i.e., any two are placed peacefully) and the placing is also invariant to a 180-degree rotation.',\n",
       "  'problem_answer': '2^floor(n/2)*floor(n/2)!'},\n",
       " {'ids': 1,\n",
       "  'problem_identifier': 'P_Combinatorics_7',\n",
       "  'problem_text': 'Consider a row of 11 seats. A child sits on each. Each child may move by at most one seat. How many possible rearrangements are there (including the original one)?',\n",
       "  'problem_answer': '144'},\n",
       " {'ids': 2,\n",
       "  'problem_identifier': 'P_Combinatorics_8',\n",
       "  'problem_text': 'How many strings of length 5 and consisting solely of the digits 1, 2, 3 and 4 have the property that any 1 and 2 are never neighbors (i.e., the string does not contain \"12\" or \"21\" as substrings).',\n",
       "  'problem_answer': '634'},\n",
       " {'ids': 3,\n",
       "  'problem_identifier': 'P_Combinatorics_9',\n",
       "  'problem_text': 'Consider all 2^n−1 nonempty subsets of the set {1, 2,..., n}. For every such subset, we find the product of the reciprocals of each of its elements. Find the sum of all these products, as an expression of n.',\n",
       "  'problem_answer': 'n'},\n",
       " {'ids': 4,\n",
       "  'problem_identifier': 'P_Combinatorics_12',\n",
       "  'problem_text': 'Along a one-way street there are n parking lots. One-by-one n cars numbered 1 to n enter the street. Each driver i heads to his favorite parking lot a_i, and, if it is free, he occupies it. Otherwise, he continues to the next free lot and occupies it. But if all succeeding lots are occupied, he leaves for good. How many sequences {a_1, ..., a_n} are such that every driver can park, as an expression of n?',\n",
       "  'problem_answer': '(n+1)^(n-1)'},\n",
       " {'ids': 5,\n",
       "  'problem_identifier': 'P_Combinatorics_13',\n",
       "  'problem_text': 'Of 3n+1 objects, n are indistinguishable, and the remaining ones are distinct. In how many ways can we choose n objects, as an expression of n?',\n",
       "  'problem_answer': '4^n'},\n",
       " {'ids': 6,\n",
       "  'problem_identifier': 'P_Combinatorics_15',\n",
       "  'problem_text': 'Find the number of ways to fill a 2 x 21 rectangle with 2 x 1 tiles such that the tiling is same when the rectangle is flipped around the middle column.',\n",
       "  'problem_answer': '89'},\n",
       " {'ids': 7,\n",
       "  'problem_identifier': 'P_Combinatorics_16',\n",
       "  'problem_text': 'Let A be a string of n binary digits (leading 0s are allowed). How many such strings have the substring \"01\" appearing exactly m times, as an expression of m and n?',\n",
       "  'problem_answer': 'C(n+1, 2m+1)'},\n",
       " {'ids': 8,\n",
       "  'problem_identifier': 'P_Combinatorics_19',\n",
       "  'problem_text': 'Find a closed-form formula of sum_(k=1)^n C(n, k)*k^2.',\n",
       "  'problem_answer': 'n(n+1)*2^(n-2)'},\n",
       " {'ids': 9,\n",
       "  'problem_identifier': 'P_Combinatorics_20',\n",
       "  'problem_text': 'Find the number of ways to fill a 2 x 5 rectangle with 1 x 1 tiles and three-cell L-shaped tiles (i.e., a 2 x 2 tile without a corner).',\n",
       "  'problem_answer': '87'},\n",
       " {'ids': 10,\n",
       "  'problem_identifier': 'P_Combinatorics_22',\n",
       "  'problem_text': 'Let 1≤r≤n and consider all subsets of r elements of the set {1, 2, ..., n}. What is the arithmetic mean of the smallest element of these subsets, as an expression of n and r?',\n",
       "  'problem_answer': '(n+1)/(r+1)'},\n",
       " {'ids': 11,\n",
       "  'problem_identifier': 'P_Combinatorics_23',\n",
       "  'problem_text': 'Among all strings of length n consisting of digit 1, 2, 3 and 4, how many of them have an even number of 1s, as an expression of n?',\n",
       "  'problem_answer': '(4^n+2^n)/2'},\n",
       " {'ids': 12,\n",
       "  'problem_identifier': 'P_Combinatorics_24',\n",
       "  'problem_text': 'Is it possible to label the edges of a cube by 1, 2,..., 12 so that, at each vertex, the labels of the edges leaving that vertex have the same sum?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 13,\n",
       "  'problem_identifier': 'P_Combinatorics_33',\n",
       "  'problem_text': 'From the set {1, 2, 3, ..., n+1}, select quadruple (x, y, z, u), with duplicate allowed, such that u>max(x, y, z). In how many ways can this be done?',\n",
       "  'problem_answer': 'sum_(i=1)^n i^3, or equivalently (n(n+1)/2)^2'},\n",
       " {'ids': 14,\n",
       "  'problem_identifier': 'P_Combinatorics_34',\n",
       "  'problem_text': 'Among all sequences of positive integer numbers have sum n, for integer k<n-1, how many times does the number k appear, as an expression of n and k?',\n",
       "  'problem_answer': '(n-k+3)*2^(n-k-2)'},\n",
       " {'ids': 15,\n",
       "  'problem_identifier': 'P_Combinatorics_36',\n",
       "  'problem_text': 'n people sit around a circular table. How many of the n! arrangements are distinct (i.e., do not have the same neighboring relations), as an expression of n?',\n",
       "  'problem_answer': 'n!/(2n)'},\n",
       " {'ids': 16,\n",
       "  'problem_identifier': 'P_Inequality_6',\n",
       "  'problem_text': 'Let A=(1/2)*(3/4)*(5/6)*...*(99/100). Is A greater than 0.1?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 17,\n",
       "  'problem_identifier': 'P_Inequality_7',\n",
       "  'problem_text': 'If a, b, c are side lengths of a (possibly degenerate) triangle, then what is the largest value of a/(b+c)+b/(a+c)+c/(a+b)?',\n",
       "  'problem_answer': '2'},\n",
       " {'ids': 18,\n",
       "  'problem_identifier': 'P_Inequality_9',\n",
       "  'problem_text': \"Let x(n)=2^2^2^...^2 with n 2's, and y(n)=3^3^3^...^3 with n 3's. Among integer values of n≥3, what is the relationship between x(n) and y(n-1): x(n)>y(n-1) for all n≥3, x(n)<y(n-1) for all n≥3, or x(n)>y(n-1) for some n and x(n)<y(n-1) for some other n?\",\n",
       "  'problem_answer': 'x(n)<y(n-1) for all n≥3'},\n",
       " {'ids': 19,\n",
       "  'problem_identifier': 'P_Inequality_12',\n",
       "  'problem_text': 'Let (a_1, a_2, ..., a_n) be a permutation of (1, 2, ..., n). What is the smallest value of sum_(k=1)^n a_k/k^2 as an expression of n?',\n",
       "  'problem_answer': 'sum_(k=1)^n 1/k'},\n",
       " {'ids': 20,\n",
       "  'problem_identifier': 'P_Inequality_13',\n",
       "  'problem_text': 'For real numbers a, b, c satisfying a+b+c=9, what is the smallest value of a^2+b^2+c^2-(ab+bc+ac)-(a+b+c)?',\n",
       "  'problem_answer': '-9'},\n",
       " {'ids': 21,\n",
       "  'problem_identifier': 'P_Inequality_15',\n",
       "  'problem_text': 'Find the minimum value of a^4+b^4+c^4-a^2*bc-b^2*ac-c^2*ab for positive numbers a, b, c.',\n",
       "  'problem_answer': '0'},\n",
       " {'ids': 22,\n",
       "  'problem_identifier': 'P_Inequality_17',\n",
       "  'problem_text': 'Let x_1, x_2, ..., x_n>0 and x_1+x_2+...+x_n=1. Let s be the greatest of x_1/(1+x_1), x_2/(1+x_1+x_2), ..., x_n/(1+x_1+x_2+...+x_n). What is the smallest value of s as an expression of n?',\n",
       "  'problem_answer': '1-1/2^(1/n)'},\n",
       " {'ids': 23,\n",
       "  'problem_identifier': 'P_Inequality_20',\n",
       "  'problem_text': 'For real values a, b, c, with a^2+b^2+c^2=1, what is the difference between the largest and smallest possible values of ab+bc+ac?',\n",
       "  'problem_answer': '3/2'},\n",
       " {'ids': 24,\n",
       "  'problem_identifier': 'P_Inequality_21',\n",
       "  'problem_text': 'For positive real number a, b and positive integer n, what is the largest value of (a*b^n)^(1/(n+1))/(a+bn) as an expression of n?',\n",
       "  'problem_answer': '1/(n+1)'},\n",
       " {'ids': 25,\n",
       "  'problem_identifier': 'P_Inequality_22',\n",
       "  'problem_text': 'Let a, b, c be the lengths of the three sides of a (possibly degenerate) triangle. What is the largest value of (a^2+b^2+c^2)/(ab+bc+ac)?',\n",
       "  'problem_answer': '2'},\n",
       " {'ids': 26,\n",
       "  'problem_identifier': 'P_Inequality_23',\n",
       "  'problem_text': 'Let a, b, c>0. What is the minimum value of abc-(a+b-c)(a+c-b)(b+c-a)?',\n",
       "  'problem_answer': '0'},\n",
       " {'ids': 27,\n",
       "  'problem_identifier': 'P_Inequality_25',\n",
       "  'problem_text': 'For how many values of n in {1, 2, ..., 100}, do we have (1/2)*(3/4)*(5/6)*...*(2n-1)/(2n)≤1/sqrt(3n+1)?',\n",
       "  'problem_answer': '100 (i.e., all values)'},\n",
       " {'ids': 28,\n",
       "  'problem_identifier': 'P_Inequality_28',\n",
       "  'problem_text': 'Let A=(1/2)*(3/4)*(5/6)*...*(99/100). Is A greater than 1/15?',\n",
       "  'problem_answer': 'Yes'},\n",
       " {'ids': 29,\n",
       "  'problem_identifier': 'P_Inequality_32',\n",
       "  'problem_text': \"Let x(n)=3^3^3^...^3 with n 3's, and y(n)=4^4^4^...^4 with n 4's. Among integer values of n≥2, what is the relationship between x(n) and y(n-1): x(n)>y(n-1) for all n≥2, x(n)<y(n-1) for all n≥2, or x(n)>y(n-1) for some n and x(n)<y(n-1) for some other n?\",\n",
       "  'problem_answer': 'x(n)>y(n-1) for all n≥2'},\n",
       " {'ids': 30,\n",
       "  'problem_identifier': 'P_Inequality_35',\n",
       "  'problem_text': 'For how many values of n in {101, ..., 1000} is 1/(n+1)+1/(n+2)+...+1/(2n)<3/4?',\n",
       "  'problem_answer': '900 (i.e., all values)'},\n",
       " {'ids': 31,\n",
       "  'problem_identifier': 'P_Inequality_37',\n",
       "  'problem_text': 'For positive a, b, c, what is the smallest value of sqrt(ab+bc+ac)/(abc)^(1/3)?',\n",
       "  'problem_answer': 'sqrt(3)'},\n",
       " {'ids': 32,\n",
       "  'problem_identifier': 'P_Inequality_39',\n",
       "  'problem_text': 'For how many values of n in {101, ..., 1000} is 1/(n+1)+1/(n+2)+...+1/(2n)>1/2?',\n",
       "  'problem_answer': '900 (i.e., all values)'},\n",
       " {'ids': 33,\n",
       "  'problem_identifier': 'P_Inequality_43',\n",
       "  'problem_text': 'Let a, b, c>0. What is the largest value of ((a+b+c)/(abc))/(1/a^2+1/b^2+1/c^2)?',\n",
       "  'problem_answer': '1'},\n",
       " {'ids': 34,\n",
       "  'problem_identifier': 'P_Inequality_45',\n",
       "  'problem_text': 'For a, b, c>0, are there possible values of them to make (a^3*b+b^3*c+c^3*a)/(a+b+c) less than abc?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 35,\n",
       "  'problem_identifier': 'P_Inequality_48',\n",
       "  'problem_text': 'For how many values of n in {1, 2, ..., 100} do we have (n+1)^n≥2^n*n!?',\n",
       "  'problem_answer': '100 (i.e., all values)'},\n",
       " {'ids': 36,\n",
       "  'problem_identifier': 'P_Number-Theory_1',\n",
       "  'problem_text': 'Are there integer solutions to the equation (x^2-1)(y^2-1)+1985=z^2?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 37,\n",
       "  'problem_identifier': 'P_Number-Theory_2',\n",
       "  'problem_text': 'Do there exist positive integers x, y such that x+y, x+2y and 2x+y are both squares?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 38,\n",
       "  'problem_identifier': 'P_Number-Theory_14',\n",
       "  'problem_text': 'Is 4^545+545^4 a prime number?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 39,\n",
       "  'problem_identifier': 'P_Number-Theory_15',\n",
       "  'problem_text': 'Can there exist a polynomial f(x) of integer coefficients with f(7)=11 and f(11)=13?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 40,\n",
       "  'problem_identifier': 'P_Number-Theory_18',\n",
       "  'problem_text': 'Find all integer solutions to the equation x^2-3y^2=17.',\n",
       "  'problem_answer': 'No integer solutions'},\n",
       " {'ids': 41,\n",
       "  'problem_identifier': 'P_Number-Theory_21',\n",
       "  'problem_text': 'Find all integer solutions to the equation 19x^3-84y^2=1984.',\n",
       "  'problem_answer': 'No integer solutions'},\n",
       " {'ids': 42,\n",
       "  'problem_identifier': 'P_Number-Theory_23',\n",
       "  'problem_text': 'Find all positive integer solutions to the equation x^3+3=4y(y+1).',\n",
       "  'problem_answer': 'No positive integer solutions'},\n",
       " {'ids': 43,\n",
       "  'problem_identifier': 'P_Number-Theory_27',\n",
       "  'problem_text': 'Let n be 22...22, with a total of 1980 digits. What is n mod 1982?',\n",
       "  'problem_answer': '0'},\n",
       " {'ids': 44,\n",
       "  'problem_identifier': 'P_Number-Theory_28',\n",
       "  'problem_text': 'How often does the factor 2 occur in the product P(n)=(n+1)(n+2)...(2n), as an expression of n?',\n",
       "  'problem_answer': 'n times'},\n",
       " {'ids': 45,\n",
       "  'problem_identifier': 'P_Number-Theory_30',\n",
       "  'problem_text': 'Find all integer solutions to the equation 15x^2-7y^2=9.',\n",
       "  'problem_answer': 'No integer solutions'},\n",
       " {'ids': 46,\n",
       "  'problem_identifier': 'P_Number-Theory_33',\n",
       "  'problem_text': 'Find the smallest positive integer a, so that 1971 | 50^n+23^n*a for odd n.',\n",
       "  'problem_answer': '512'},\n",
       " {'ids': 47,\n",
       "  'problem_identifier': 'P_Number-Theory_34',\n",
       "  'problem_text': 'Are there integer solutions to the equation (x^2-1)(y^2-1)+1984=z^2?',\n",
       "  'problem_answer': 'Yes'},\n",
       " {'ids': 48,\n",
       "  'problem_identifier': 'P_Number-Theory_35',\n",
       "  'problem_text': 'A sequence starts with 1, 9, 7, 7, 4, 7, 5, 3, 9, 4, 1... From the fifth number, every number is the sum of the preceding four numbers mod 10. Could the sequence contain 3, 2, 6, 9 as a subsequence somewhere?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 49,\n",
       "  'problem_identifier': 'P_Number-Theory_36',\n",
       "  'problem_text': 'Find all possible values of x^2 mod 30 for prime number x>30.',\n",
       "  'problem_answer': '1 and 19'},\n",
       " {'ids': 50,\n",
       "  'problem_identifier': 'P_Number-Theory_39',\n",
       "  'problem_text': 'Find all non-negative integer solutions to the equation x^3+8x^2-6x+8=y^3.',\n",
       "  'problem_answer': 'x=0, y=2 and x=9, y=11'},\n",
       " {'ids': 51,\n",
       "  'problem_identifier': 'P_Number-Theory_41',\n",
       "  'problem_text': 'Find all possible values of n such that n^2-19n+89 is a perfect square and n>11.',\n",
       "  'problem_answer': 'No such values'},\n",
       " {'ids': 52,\n",
       "  'problem_identifier': 'P_Number-Theory_46',\n",
       "  'problem_text': 'For positive integer n for which 2n+1 and 3n+1 are both perfect squares, what are possible values of n mod 40?',\n",
       "  'problem_answer': '0 is the only possible value'},\n",
       " {'ids': 53,\n",
       "  'problem_identifier': 'P_Number-Theory_47',\n",
       "  'problem_text': 'Find all integer values of m such that 1000^m-1 is a factor of 1978^m-1.',\n",
       "  'problem_answer': 'No such values'},\n",
       " {'ids': 54,\n",
       "  'problem_identifier': 'P_Number-Theory_54',\n",
       "  'problem_text': 'How many pairs of positive odd integer (a, b) both under 100 make a^2+b^2 be a perfect square?',\n",
       "  'problem_answer': '0 pairs'},\n",
       " {'ids': 55,\n",
       "  'problem_identifier': 'P_Number-Theory_57',\n",
       "  'problem_text': 'Find all possible values of 4^n+15n-1 mod 9 for interger n≥0.',\n",
       "  'problem_answer': '0 is the only possible value'},\n",
       " {'ids': 56,\n",
       "  'problem_identifier': 'P_Number-Theory_58',\n",
       "  'problem_text': 'Find all possible values of 2(1^k+2^k+...+n^k) mod n(n+1) for odd k.',\n",
       "  'problem_answer': '0 is the only possible value'},\n",
       " {'ids': 57,\n",
       "  'problem_identifier': 'P_Number-Theory_61',\n",
       "  'problem_text': 'For positive integers a, b, c, d, n satisfying ab=cd, can a^n+b^n+c^n+d^n be a prime number?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 58,\n",
       "  'problem_identifier': 'P_Number-Theory_62',\n",
       "  'problem_text': 'Find all prime number values of n>2 such that 1+2+...+(n-2)+(n-1) is a factor of 1*2*...*(n-2)*(n-1).',\n",
       "  'problem_answer': 'No such values'},\n",
       " {'ids': 59,\n",
       "  'problem_identifier': 'P_Number-Theory_64',\n",
       "  'problem_text': 'Find all positive integer solutions to the equation m^2+(m+1)^2=n^4+(n+1)^4.',\n",
       "  'problem_answer': 'No positive integer solutions'},\n",
       " {'ids': 60,\n",
       "  'problem_identifier': 'P_Number-Theory_68',\n",
       "  'problem_text': 'Find all integer solutions to the equation (n-2)^2+(n-1)^2+n^2+(n+1)^2+(n+2)^2=m^2.',\n",
       "  'problem_answer': 'No integer solutions'},\n",
       " {'ids': 61,\n",
       "  'problem_identifier': 'P_Number-Theory_70',\n",
       "  'problem_text': 'Find all positive integer solutions to the equation 2^x+1=3^y.',\n",
       "  'problem_answer': 'x=1, y=1 and x=3, y=2'},\n",
       " {'ids': 62,\n",
       "  'problem_identifier': 'P_Number-Theory_71',\n",
       "  'problem_text': 'What is the smallest value of |12^m-5^n| for positive integers m and n?',\n",
       "  'problem_answer': '7'},\n",
       " {'ids': 63,\n",
       "  'problem_identifier': 'P_Number-Theory_72',\n",
       "  'problem_text': 'If positive integers x and y are such that x^2+2y^2 is a prime number, what are all possible values of x^2+2y^2 mod 8?',\n",
       "  'problem_answer': '1 and 3'},\n",
       " {'ids': 64,\n",
       "  'problem_identifier': 'P_Number-Theory_75',\n",
       "  'problem_text': 'Are there integer solutions to the equation (x^2-1)(y^2-1)+1981=z^2?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 65,\n",
       "  'problem_identifier': 'P_Number-Theory_77',\n",
       "  'problem_text': 'If the integer m ends with digit 5, what are possible values of 12^m+9^m+8^m+6^m mod 1991?',\n",
       "  'problem_answer': '0 is the only possible value'},\n",
       " {'ids': 66,\n",
       "  'problem_identifier': 'P_Number-Theory_79',\n",
       "  'problem_text': 'Let integer p, q satisfy p/q=1-1/2+1/3-1/4+...-1/1318+1/1319 with gcd(p,q)=1. What is p mod 1979?',\n",
       "  'problem_answer': '0'},\n",
       " {'ids': 67,\n",
       "  'problem_identifier': 'P_Polynomial_2',\n",
       "  'problem_text': 'Solve the equation x^4+a^4-3a*x^3+3a^3*x=0 for a≠0 as an expression of a.',\n",
       "  'problem_answer': 'a(1±sqrt(2)), a(1±sqrt(5))/2'},\n",
       " {'ids': 68,\n",
       "  'problem_identifier': 'P_Polynomial_3',\n",
       "  'problem_text': 'For a polynomial f(x) of degree n with integer coefficients, if both f(0) and f(1) are odd, at most how many integer-valued zeros can it have?',\n",
       "  'problem_answer': 'At most 0 integer-valued zeros (i.e., no integer zeros)'},\n",
       " {'ids': 69,\n",
       "  'problem_identifier': 'P_Polynomial_5',\n",
       "  'problem_text': 'Find all integer solutions to the equation x^3+x^2*y+x*y^2+y^3=8(x^2+xy+y^2+1).',\n",
       "  'problem_answer': '(x=8, y=2), (x=2, y=8)'},\n",
       " {'ids': 70,\n",
       "  'problem_identifier': 'P_Polynomial_11',\n",
       "  'problem_text': 'If x_1, x_2 are the two roots of the polynomial x^2-6x+1, for how many integer n in {61, 62, ..., 120} does x_1^n+x_2^n divided by 5 give a remainder of 4?',\n",
       "  'problem_answer': '20'},\n",
       " {'ids': 71,\n",
       "  'problem_identifier': 'P_Polynomial_14',\n",
       "  'problem_text': 'Factorize x^4+x^2+1 into terms of integer coefficients.',\n",
       "  'problem_answer': '(x^2+x+1)(x^2-x+1)'},\n",
       " {'ids': 72,\n",
       "  'problem_identifier': 'P_Polynomial_15',\n",
       "  'problem_text': 'Find all possible values of f(x) at x≠0 such that xf(y)+yf(x)=(x+y)f(x)f(y) for all x, y.',\n",
       "  'problem_answer': 'f(x)=0 or f(x)=1'},\n",
       " {'ids': 73,\n",
       "  'problem_identifier': 'P_Polynomial_16',\n",
       "  'problem_text': 'The polynomial f(x) has integer coefficients and there exists an integer k such that f(k), f(k+1), f(k+2) are all divisible by 3. For integer m in {k+3, k+4, ..., 2k}, at least how many of them have f(m) divisible by 3, as an expression of k?',\n",
       "  'problem_answer': 'k-2 (i.e., all of them)'},\n",
       " {'ids': 74,\n",
       "  'problem_identifier': 'P_Polynomial_17',\n",
       "  'problem_text': 'Let f(x) be a polynomial of degree n and f(x)=x has no real number solutions. What is the maximum number of real number solutions does f(f(x))=x have?',\n",
       "  'problem_answer': '0'},\n",
       " {'ids': 75,\n",
       "  'problem_identifier': 'P_Polynomial_18',\n",
       "  'problem_text': 'Factorize x^8+x^4+1 into terms of integer coefficients.',\n",
       "  'problem_answer': '(x^2+x+1)(x^2-x+1)(x^4-x^2+1)'},\n",
       " {'ids': 76,\n",
       "  'problem_identifier': 'P_Polynomial_20',\n",
       "  'problem_text': 'Find a polynomial f(x) such that x*f(x−1)=(x+1)*f(x) for all x?',\n",
       "  'problem_answer': 'f(x)=0'},\n",
       " {'ids': 77,\n",
       "  'problem_identifier': 'P_Polynomial_22',\n",
       "  'problem_text': 'Find all real-valued solutions to the equation x^8+4x^6−10x^4+4x^2+1=0.',\n",
       "  'problem_answer': 'x=±1'},\n",
       " {'ids': 78,\n",
       "  'problem_identifier': 'P_Polynomial_24',\n",
       "  'problem_text': 'Find the smallest value of the polynomial f(x)=x^3(x^3+1)(x^3+2)(x^3+3).',\n",
       "  'problem_answer': '-1'},\n",
       " {'ids': 79,\n",
       "  'problem_identifier': 'P_Polynomial_26',\n",
       "  'problem_text': 'For what value(s) of k is x^3+y^3+z^3+kxyz divisible by x+y+z?',\n",
       "  'problem_answer': '-3 is the only possible value'},\n",
       " {'ids': 80,\n",
       "  'problem_identifier': 'P_Polynomial_27',\n",
       "  'problem_text': 'Let f(x) be a polynomial with integer coefficients, with f(x)=5 at four distinct integers x=a, b, c, d. If f(x) has degree n, then at most how many integer values k are there with f(k)=8?',\n",
       "  'problem_answer': 'At most 0 values (i.e., no such values)'},\n",
       " {'ids': 81,\n",
       "  'problem_identifier': 'P_Polynomial_28',\n",
       "  'problem_text': 'Let P(x) be a polynomial of degree n such that P(k)=k/(k+1) for k=0, 1, ..., n. Find P(n+1).',\n",
       "  'problem_answer': '((-1)^(n+1)+n+1)/(n+2) (i.e., 1 if for odd n, and n/(n+2) for even n)'},\n",
       " {'ids': 82,\n",
       "  'problem_identifier': 'P_Polynomial_29',\n",
       "  'problem_text': 'Factorize x^9+x^4-x-1 into terms of integer coefficients.',\n",
       "  'problem_answer': '(x+1)(x-1)(x^2+1)(x^3−x^2+1)(x^2+x+1)'},\n",
       " {'ids': 83,\n",
       "  'problem_identifier': 'P_Polynomial_31',\n",
       "  'problem_text': 'Suppose that for x, y we have x+y=a and x^3+y^3=b. Find x^2+y^2, as an expression of a and b.',\n",
       "  'problem_answer': '(a^3+2b)/(3a)'},\n",
       " {'ids': 84,\n",
       "  'problem_identifier': 'P_Polynomial_35',\n",
       "  'problem_text': 'Let f(x)=(x^1958+x^1957+2)^1959=a_0+a_1*x+···+a_n*x^n. Find a_0−a_1/2−a_2/2+a_3−a_4/2−a_5/2+a_6-a_7/2-a_8/2+....',\n",
       "  'problem_answer': '1'},\n",
       " {'ids': 85,\n",
       "  'problem_identifier': 'P_Polynomial_36',\n",
       "  'problem_text': 'For the polynomial x^3+4x^2+2x+a, what value of a results in one root being equal to the sum of the other two roots?',\n",
       "  'problem_answer': '-4'},\n",
       " {'ids': 86,\n",
       "  'problem_identifier': 'P_Polynomial_38',\n",
       "  'problem_text': 'Find all real-valued solutions to the equation 4x^11+4x^10−21x^9-21x^8+17x^7+17x^6+17x^5+17x^4−21x^3-21x^2+4x+4=0.',\n",
       "  'problem_answer': 'x=±1, ±2, ±1/2'},\n",
       " {'ids': 87,\n",
       "  'problem_identifier': 'P_Polynomial_39',\n",
       "  'problem_text': 'Find the remainder of x^1959−1 divided by (x^2+1)(x^2+x+1).',\n",
       "  'problem_answer': 'x^3-1'},\n",
       " {'ids': 88,\n",
       "  'problem_identifier': 'P_Polynomial_41',\n",
       "  'problem_text': 'What is the largest integer n for which n^3+100 is divisible by n+10?',\n",
       "  'problem_answer': '890'},\n",
       " {'ids': 89,\n",
       "  'problem_identifier': 'P_Polynomial_42',\n",
       "  'problem_text': 'Let f(x) be a polynomial of degree n with integer coefficients. If there are three different integers a, b, c, such that f(a)=f(b)=f(c)=-1, then at most how many integer-valued roots can this polynomial have?',\n",
       "  'problem_answer': 'At most 0 integer-valued roots (i.e., no integer roots)'},\n",
       " {'ids': 90,\n",
       "  'problem_identifier': 'P_Polynomial_45',\n",
       "  'problem_text': 'The polynomial ax^3+bx^2+cx+d has integer coefficients a, b, c, d with ad odd and bc even. Can some values of a, b, c, d result in all three roots of the polynomial being integers?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 91,\n",
       "  'problem_identifier': 'P_Polynomial_46',\n",
       "  'problem_text': 'Find a monic polynomial (i.e. a polynomial with the leading coefficient being 1) p(x) of degree greater than 1 that satisfies x*p(x-1)=(x-2023)*p(x).',\n",
       "  'problem_answer': 'x(x-1)(x-2)...(x-2022)'},\n",
       " {'ids': 92,\n",
       "  'problem_identifier': 'P_Polynomial_47',\n",
       "  'problem_text': 'The polynomial 1-x+x^2-x^3+...+x^8-x^9 may be written in the form a_0+a_1*y+a_2*y^2+...+a_9*y^9, where y=x+1 and each a_i is a constant. Find the value of a_2.',\n",
       "  'problem_answer': '120'},\n",
       " {'ids': 93,\n",
       "  'problem_identifier': 'P_Polynomial_48',\n",
       "  'problem_text': 'Let x_1, x_2, x_3 be the roots of x^3+3x^2-7x+1. Find x_1^2+x_2^2+x_3^2.',\n",
       "  'problem_answer': '23'},\n",
       " {'ids': 94,\n",
       "  'problem_identifier': 'P_Polynomial_49',\n",
       "  'problem_text': 'Is it possible that each of the polynomials P(x)=ax^2+bx+c, Q(x)=cx^2+ax+b and R(x)=bx^2+cx+a has two different real roots with positive values of a, b, c?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 95,\n",
       "  'problem_identifier': 'P_Sequence_1',\n",
       "  'problem_text': 'Let binary string w_1=0, and w_(n+1) is generated by replacing each 0 in w_n by 001, and each 1 in w_n by 0. Thus, we have w_2=001, w_3=0010010, etc. Does this operation result in a periodic binary string in the infinite limit, and if so, what is its period?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 96,\n",
       "  'problem_identifier': 'P_Sequence_2',\n",
       "  'problem_text': 'Let {x_n}, {y_n}, {z_n} be three sequences with positive initial terms x_1, y_1, z_1, defined as x_(n+1)=y_n+1/z_n, y_(n+1)=z_n+1/x_n, z_(n+1)=x_n+1/y_n. Let w_n be the maximum value of x_n, y_n, z_n. For different values of x_1, y_1, z_1, do we have w_200 always greater than 20, always smaller than 20, or sometimes greater and sometimes smaller than 20?',\n",
       "  'problem_answer': 'Always greater than 20'},\n",
       " {'ids': 97,\n",
       "  'problem_identifier': 'P_Sequence_4',\n",
       "  'problem_text': 'Let x_1, x_2<100 be two positive integers. For k≥3, let x_k be the smallest of the absolute values of the pairwise differences of the preceding terms. What are possible values of x_12?',\n",
       "  'problem_answer': '0 is the only possible value'},\n",
       " {'ids': 98,\n",
       "  'problem_identifier': 'P_Sequence_8',\n",
       "  'problem_text': 'Find the limit of the expression 1/(1*2*3*4)+1/(2*3*4*5)+...+1/(n(n+1)(n+2)(n+3)) as n→∞.',\n",
       "  'problem_answer': '1/18'},\n",
       " {'ids': 99,\n",
       "  'problem_identifier': 'P_Sequence_11',\n",
       "  'problem_text': 'Define a sequence with x_1=1/2, x_(k+1)=x_k^2+x_k. What is the integer part of the sum 1/(x_1+1)+1/(x_2+1)+1/(x_3+1)+...+1/(x_100+1)?',\n",
       "  'problem_answer': '1'},\n",
       " {'ids': 100,\n",
       "  'problem_identifier': 'P_Sequence_13',\n",
       "  'problem_text': 'A sequence of positive integers a_0, a_1, ..., a_100 is defined as a_1>a_0, a_n=3a_(n-1)-2a_(n-2). Is it possible to have a_100<2^99 for certain such sequences, and if so, find the sequence with the largest sum a_0+a_1+...+a_100.',\n",
       "  'problem_answer': 'Such a sequence is not possible'},\n",
       " {'ids': 101,\n",
       "  'problem_identifier': 'P_Sequence_15',\n",
       "  'problem_text': 'Does there exist a positive sequence {a_n}, such that both sum a_n and sum 1/(n^2*a_n) converge to finite values?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 102,\n",
       "  'problem_identifier': 'P_Sequence_23',\n",
       "  'problem_text': 'A sequence is defined as a_1=0, |a_n|=|a_(n-1)+1|. What is the smallest value of (a_1+a_2+...+a_n)/n for even n?',\n",
       "  'problem_answer': '-1/2'},\n",
       " {'ids': 103,\n",
       "  'problem_identifier': 'P_Sequence_29',\n",
       "  'problem_text': 'Does the limit exist for the sequence defined as a_n=((2^3-1)/(2^3+1))*((3^3-1)/(3^3+1))*((4^3-1)/(4^3+1))*...*((n^3-1)/(n^3+1)), and if so, find it.',\n",
       "  'problem_answer': 'The limit exists and is equal to 2/3'},\n",
       " {'ids': 104,\n",
       "  'problem_identifier': 'P_Sequence_30',\n",
       "  'problem_text': 'For a≤b, define the sequence as a_1=a, a_2=b, a_(n+2)=(a_(n+1)+a_n)/2. Does the limit exist for this sequence, and if so, find it as an expression in terms of a and b.',\n",
       "  'problem_answer': 'The limit exists and is equal to (a+3b)/2'},\n",
       " {'ids': 105,\n",
       "  'problem_identifier': 'P_Sequence_31',\n",
       "  'problem_text': 'What are possible values of the limit of the sequence a_n=sqrt(a_(n-1))+sqrt(a_(n-2)), for different initial values of 0<a_0, a_1<1, possibly as an expression of a_0, a_1?',\n",
       "  'problem_answer': 'The only possible value for the limit is 4'},\n",
       " {'ids': 106,\n",
       "  'problem_identifier': 'P_Sequence_32',\n",
       "  'problem_text': 'Given a set of n positive numbers a_1≤a_2≤...≤a_n, with the sum of the n(n-1)/2 pairwise products being equal to 1. Is it possible that a_1+a_2+...+a_(n-1) is greater than or equal to sqrt(2)?',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 107,\n",
       "  'problem_identifier': 'P_Sequence_34',\n",
       "  'problem_text': 'The positive real numbers x_0, x_1, ..., x_1995 satisfy x_0=x_1995 and x_i+2/x_i=2x_(i+1)+1/x_(i+1) for i=0, ..., 1994. Find the maximum value that x_0 can have.',\n",
       "  'problem_answer': '2^997'},\n",
       " {'ids': 108,\n",
       "  'problem_identifier': 'P_Sequence_37',\n",
       "  'problem_text': 'The sequence {x_n} is defined by x_1=2, x_(n+1)=(2+x_n)/(1−2x_n). Must x_n≠0 for all n in the sequence?',\n",
       "  'problem_answer': 'Yes'},\n",
       " {'ids': 109,\n",
       "  'problem_identifier': 'P_Sequence_43',\n",
       "  'problem_text': 'Is it possible to select an infinite geometric sequence from the sequence 1, 1/2, 1/4, 1/8, 1/16, ... such that the sum is 1/7, and if so, what is the sequence?',\n",
       "  'problem_answer': 'Yes, and the sequence is 1/2^3, 1/2^6, 1/2^9, ...'},\n",
       " {'ids': 110,\n",
       "  'problem_identifier': 'P_Sequence_44',\n",
       "  'problem_text': 'How many sequences of positive numbers a_0, a_1, a_2, ... satisfy the conditions of a_0=1, a_(n+2)=a_n-a_(n+1) for all n≥0?',\n",
       "  'problem_answer': 'Exactly one sequence'},\n",
       " {'ids': 111,\n",
       "  'problem_identifier': 'P_Sequence_45',\n",
       "  'problem_text': 'A sequence {a_n} is defined by a_1=1, a_(n+1)=a_n+1/a_n^2. Do we have a_9000>30?',\n",
       "  'problem_answer': 'Yes'},\n",
       " {'ids': 112,\n",
       "  'problem_identifier': 'P_Sequence_46',\n",
       "  'problem_text': 'In how many ways can 12 geometric progressions be selected from the numbers 1, 2, ..., 100 such that every number is in exactly one progression?',\n",
       "  'problem_answer': '0 (no such selection is possible)'},\n",
       " {'ids': 113,\n",
       "  'problem_identifier': 'P_Sequence_48',\n",
       "  'problem_text': 'A function f(x) satisfies f(x+1)+f(x-1)=sqrt(2)*f(x). Determine whether the function is periodic, and if so, find its period.',\n",
       "  'problem_answer': 'The function is periodic, with a period of 8'},\n",
       " {'ids': 114,\n",
       "  'problem_identifier': 'P_Sequence_49',\n",
       "  'problem_text': 'The sequence {x_n} is defined by x_1=2, x_(n+1)=(2+x_n)/(1−2x_n), and it is known that there are no x_n=0. Could x_p=x_q for any p≠q, and if so, find the smallest possible value of p+q.',\n",
       "  'problem_answer': 'No'},\n",
       " {'ids': 115,\n",
       "  'problem_identifier': 'P_Sequence_50',\n",
       "  'problem_text': 'Let a sequence be the defined as a_1=a_2=1, and a_n=(a_(n-1)^2+2)/a_(n-2). How many values in a_1, a_2, ..., a_100 are integers?',\n",
       "  'problem_answer': '100 (i.e., all values)'}]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_string_by_(string):\n",
    "    return string.split('_')[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Number-Theory'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split_string_by_(dataset[60]['problem_identifier'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "\n",
    "question = dataset[i]['problem_text']\n",
    "type = split_string_by_(dataset[i]['problem_identifier'])\n",
    "gold_answer = dataset[i]['problem_answer']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'On a chess board, two rooks are placed peacefully if they are not on the same row or column. For an n x n chess board, find the number of ways that n rooks can be placed peacefully (i.e., any two are placed peacefully) and the placing is also invariant to a 180-degree rotation.'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "任务分解 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total tokens: 471\n",
      "['How many ways can the first rook be placed on the chess board?', 'How many ways can the second rook be placed on the chess board, ensuring it is peacefully placed with the first rook?', 'How many ways can the third rook be placed on the chess board, ensuring it is peacefully placed with the first two rooks?', 'Continue this process until all n rooks have been placed on the chess board.', 'Determine how this placement is invariant to a 180-degree rotation.', 'Calculate the total number of ways n rooks can be placed peacefully on an n x n chess board.']\n",
      "{0: 'How many ways can the first rook be placed on the chess board?', 1: 'How many ways can the second rook be placed on the chess board, ensuring it is peacefully placed with the first rook?', 2: 'How many ways can the third rook be placed on the chess board, ensuring it is peacefully placed with the first two rooks?', 3: 'Continue this process until all n rooks have been placed on the chess board.', 4: 'Determine how this placement is invariant to a 180-degree rotation.', 5: 'Calculate the total number of ways n rooks can be placed peacefully on an n x n chess board.'}\n"
     ]
    }
   ],
   "source": [
    "def decompose_sql(question, type):   \n",
    "\n",
    "    prompt_for_decompose = f\"\"\"\n",
    "\n",
    "I will now give you a math problem. The type of problem is {type}. Please help me translate this math problem into a series of step-by-step sub-problems.\n",
    "\n",
    "1 examples are as follows:\n",
    "Question: Four years ago, Kody was only half as old as Mohamed. If Mohamed is currently twice 30 years old, how old is Kody? \n",
    "Answer: \n",
    "To solve the question \"How old is Kody?\", we need to solve the following problems step by step:\n",
    "1. How old is Mohamed now?\n",
    "2. How old was Mohamed four years ago?\n",
    "3. How old was Kody four years ago?\n",
    "4. How old is Kody now?\n",
    "\n",
    "Now the command is {question}, please decompose it into a series of easy-to-solve steps like the examples.\n",
    "Answer Format: (Please write each broken-down question step on a separate line, starting with a number.)\n",
    "To solve the question \"xxx\", we need to solve the following problems step by step:\n",
    "1. sub-question 1\n",
    "2. sub-question 2\n",
    "3. sub-question 3\n",
    "...\n",
    "\"\"\"\n",
    "\n",
    "    Q = {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": prompt_for_decompose\n",
    "    }\n",
    "    # Query = Example+[Q]\n",
    "    Query = [Q]\n",
    "    result = askChatGPT(Query, model=GPT_MODEL, temperature=1)\n",
    "    return result\n",
    "\n",
    "\n",
    "def convert_steps_to_format(decom_commands):\n",
    "    # 截取“we need to know:”后的内容\n",
    "    start_index = decom_commands.find(\"we need to solve the following problems step by step:\") + len(\"we need to solve the following problems step by step:\")\n",
    "    subtasks_text = decom_commands[start_index:].strip()\n",
    "    # 将每个子任务单独列出\n",
    "    subtasks = subtasks_text.split('\\n')\n",
    "    subtasks = [task.strip().split('. ', 1)[-1] for task in subtasks]\n",
    "    steps_dict = {index: value for index, value in enumerate(subtasks)}\n",
    "    return subtasks, steps_dict\n",
    "\n",
    "decompose_steps = decompose_sql(question, type)\n",
    "steps, steps_dict = convert_steps_to_format(decompose_steps)\n",
    "num_steps = len(steps)\n",
    "print(steps)\n",
    "print(steps_dict)  # 只是加了一个问题的编号而已."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6\n"
     ]
    }
   ],
   "source": [
    "print(len(steps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total tokens: 220\n",
      "total tokens: 256\n",
      "total tokens: 218\n",
      "total tokens: 254\n",
      "total tokens: 216\n",
      "total tokens: 252\n",
      "total tokens: 368\n",
      "total tokens: 404\n",
      "total tokens: 399\n",
      "total tokens: 435\n",
      "total tokens: 376\n",
      "total tokens: 412\n",
      "total tokens: 381\n",
      "total tokens: 417\n",
      "total tokens: 374\n",
      "total tokens: 410\n",
      "total tokens: 446\n",
      "total tokens: 482\n",
      "total tokens: 553\n",
      "total tokens: 589\n",
      "total tokens: 546\n",
      "total tokens: 582\n",
      "total tokens: 620\n",
      "total tokens: 656\n",
      "total tokens: 485\n",
      "total tokens: 521\n",
      "total tokens: 485\n",
      "total tokens: 521\n",
      "total tokens: 485\n",
      "total tokens: 521\n",
      "total tokens: 696\n",
      "total tokens: 732\n",
      "total tokens: 757\n",
      "total tokens: 793\n",
      "total tokens: 932\n",
      "total tokens: 968\n",
      "total tokens: 757\n",
      "total tokens: 793\n",
      "total tokens: 792\n",
      "total tokens: 828\n",
      "total tokens: 790\n",
      "total tokens: 826\n",
      "total tokens: 895\n",
      "total tokens: 931\n",
      "total tokens: 993\n",
      "total tokens: 1029\n",
      "total tokens: 862\n",
      "total tokens: 898\n",
      "total tokens: 902\n",
      "total tokens: 938\n",
      "total tokens: 1090\n",
      "total tokens: 1126\n",
      "total tokens: 999\n",
      "total tokens: 1035\n",
      "total tokens: 1053\n",
      "total tokens: 1089\n",
      "total tokens: 1063\n",
      "total tokens: 1099\n",
      "total tokens: 1057\n",
      "total tokens: 1093\n",
      "total tokens: 1103\n",
      "total tokens: 1139\n",
      "total tokens: 1103\n",
      "total tokens: 1139\n",
      "total tokens: 1098\n",
      "total tokens: 1134\n",
      "2\n",
      "(29, ['The first rook can be placed in \\\\( n^2 \\\\) ways on an \\\\( n \\\\times n \\\\) chess board.', 'Sub-problem-Id: 1; Sub-problem: How many ways can the second rook be placed on the chess board, ensuring it is peacefully placed with the first rook?; Answer: After placing the first rook, the second rook can be placed in \\\\( n-2 \\\\) ways in the same row or column as the first rook, or in \\\\( n \\\\times (n-2) \\\\) ways if it is not in the same row or column.', 'Sub-problem-Id: 2; Sub-problem: How many ways can the third rook be placed on the chess board, ensuring it is peacefully placed with the first two rooks?; Answer: After placing the first two rooks, the third rook can be placed in \\\\( (n-4) \\\\) ways if it is in the same row or column as any of the first two rooks, or in \\\\( n \\\\times (n-4) \\\\) ways if it is not in the same row or column.', 'Sub-problem-Id: 3; Sub-problem: How many ways can the fourth rook be placed on the chess board, ensuring it is peacefully placed with the first three rooks?; Answer: After placing the first three rooks, the fourth rook can be placed in \\\\( (n-6) \\\\) ways if it is in the same row or column as any of the first three rooks, or in \\\\( n \\\\times (n-6) \\\\) ways if it is not in the same row or column.', 'Sub-problem: Determine how this placement is invariant to a 180-degree rotation.\\nAnswer: In order for the placement of the rooks to be invariant to a 180-degree rotation, we need to consider the placement symmetry. For each rook placed in the first half of the chess board (e.g., first row or first column), there must be a corresponding rook placed in the second half of the chess board (e.g., last row or last column) in such a way that they are peacefully placed. This symmetry ensures that when the entire board is rotated 180 degrees, the rooks remain peacefully placed with respect to each other. Therefore, the number of ways n rooks can be placed peacefully and invariant to a 180-degree rotation is given by the product of the ways each rook can be placed while maintaining this symmetry.', 'Sub-problem: Calculate the total number of ways n rooks can be placed peacefully on an n x n chess board.\\n\\nAnswer: The total number of ways n rooks can be placed peacefully on an n x n chess board, while also being invariant to a 180-degree rotation, is given by the product of the number of ways each rook can be placed while maintaining symmetry. Therefore, the total number of ways is \\\\( n^2 \\\\times (n-2) \\\\times (n-4) \\\\times \\\\ldots \\\\) until the last step where there are two remaining options.'])\n",
      "(29, ['The first rook can be placed in \\\\( n^2 \\\\) ways on an \\\\( n \\\\times n \\\\) chess board.', 'Sub-problem-Id: 1; Sub-problem: How many ways can the second rook be placed on the chess board, ensuring it is peacefully placed with the first rook?; Answer: After placing the first rook, the second rook can be placed in \\\\( n-2 \\\\) ways in the same row or column as the first rook, or in \\\\( n \\\\times (n-2) \\\\) ways if it is not in the same row or column.', 'Sub-problem-Id: 2; Sub-problem: How many ways can the third rook be placed on the chess board, ensuring it is peacefully placed with the first two rooks?; Answer: After placing the first two rooks, the third rook can be placed in \\\\( (n-4) \\\\) ways if it is in the same row or column as any of the first two rooks, or in \\\\( n \\\\times (n-4) \\\\) ways if it is not in the same row or column.', 'Sub-problem-Id: 3; Sub-problem: How many ways can the fourth rook be placed on the chess board, ensuring it is peacefully placed with the first three rooks?; Answer: After placing the first three rooks, the fourth rook can be placed in \\\\( (n-6) \\\\) ways if it is in the same row or column as any of the first three rooks, or in \\\\( n \\\\times (n-6) \\\\) ways if it is not in the same row or column.', 'Sub-problem: Determine how this placement is invariant to a 180-degree rotation.\\nAnswer: In order for the placement of the rooks to be invariant to a 180-degree rotation, we need to consider the placement symmetry. For each rook placed in the first half of the chess board (e.g., first row or first column), there must be a corresponding rook placed in the second half of the chess board (e.g., last row or last column) in such a way that they are peacefully placed. This symmetry ensures that when the entire board is rotated 180 degrees, the rooks remain peacefully placed with respect to each other. Therefore, the number of ways n rooks can be placed peacefully and invariant to a 180-degree rotation is given by the product of the ways each rook can be placed while maintaining this symmetry.', 'Sub-problem: Calculate the total number of ways n rooks can be placed peacefully on an n x n chess board.\\nAnswer: The total number of ways n rooks can be placed peacefully on an n x n chess board, while being invariant to a 180-degree rotation, is given by the product of the ways each rook can be placed while maintaining symmetry. This can be calculated as \\\\( n^2 \\\\times (n-2) \\\\times (n-4) \\\\times ... \\\\) up to the nth rook, where each term represents the number of ways a rook can be placed without conflicting with previously placed rooks.'])\n"
     ]
    }
   ],
   "source": [
    "# TOT 求解\n",
    "N = 3  # 每个子问题进行N次proposal\n",
    "M = 2  # 通过评估选出M个最好的proposal\n",
    "\n",
    "solution = []\n",
    "\n",
    "for i in range(num_steps):\n",
    "    subtask = steps[i]\n",
    "    sys_q = f\"\"\"\n",
    "   There is a math_problem. I need you to solve it and give an answer.\n",
    "Here is the problem:\\n{question}\n",
    "\n",
    "I have broken this math problem down into a series of smaller problems. I will assign you sub-problems one by one, and provide the results of the previous sub-problems as a reference for your reasoning.\n",
    "Please solve the problem and respond according to mathematical logic.\n",
    "\"\"\"  \n",
    "\n",
    "    subask = f\"\"\"\\nThe sub-problem to solve now is xxx: {subtask}\n",
    "Based on the information above, please provide a concise and clear answer\"\"\"\n",
    "    \n",
    "    if len(solution)==0:\n",
    "        # 第一个子问题\n",
    "        query = subask\n",
    "        Q = [{'role':'system', 'content':sys_q},\n",
    "            {'role':'user', 'content':query},]\n",
    "        for n in range(N):  # 一个子问题提问N次,获取N个解\n",
    "            result = askChatGPT(Q, model=GPT_MODEL, temperature=1)\n",
    "            eval_Q = Q + [{'role':'assistant', 'content':result}]\n",
    "            eval_Q = eval_Q + [{'role':'user', 'content':\"Please provide a confidence rating for the accuracy of this solution, on a scale from 1 to 5. Only output the number.\"}]\n",
    "            score = askChatGPT(eval_Q, model=GPT_MODEL, temperature=1)\n",
    "            score = int(score)\n",
    "            \n",
    "            solution.append((score, [result]))  # 维护一整条推理路径\n",
    "            \n",
    "        solution = sorted(solution, key=lambda x: x[0])\n",
    "        solution = solution[:M]  # 剪枝\n",
    "    else:\n",
    "        temp_solution = []\n",
    "        for m in range(M):  # 因为剪枝动态维护M个推理路径\n",
    "            answersSoFar = f\"\"\"\\nSo far, the answers to the preceding sub-problems are as follows: The format is Sub-problem-Id: xxx; Sub-problem: xxx; Answer: xxx.\"\"\"\n",
    "            for index, value in enumerate(solution[m][1]):\n",
    "                try:\n",
    "                    answersSoFar += f\"\"\"\\nSub-problem-Id: {index}; Sub-problem: {steps[index]}; Answer: {value}.\"\"\"\n",
    "                except:\n",
    "                    print('warning')\n",
    "                    print(index)\n",
    "                    print(len(solution[m][1]))\n",
    "                    print(len(steps))\n",
    "                    sys.exit(0)\n",
    "            query = answersSoFar+subask\n",
    "            Q = [{'role':'system', 'content':sys_q},\n",
    "                 {'role':'user', 'content':query},]\n",
    "            for n in range(N):  # 一个子问题提问N次,获取N个解\n",
    "                result = askChatGPT(Q, model=GPT_MODEL, temperature=1)\n",
    "                eval_Q = Q + [{'role':'assistant', 'content':result}]\n",
    "                eval_Q = eval_Q + [{'role':'user', 'content':\"Please provide a confidence rating for the accuracy of this solution, on a scale from 1 to 5. Only output the number.\"}]\n",
    "                score = askChatGPT(eval_Q, model=GPT_MODEL, temperature=1)\n",
    "                score = int(score)\n",
    "                \n",
    "                temp_solution.append((solution[m][0]+score, solution[m][1]+[result]))  # 路径score累加\n",
    "        \n",
    "        # print(len(temp_solution))  # 此时temp_solution中应该有M*N种推理路径\n",
    "        solution = sorted(temp_solution, key=lambda x: x[0])\n",
    "        solution = solution[:M]  # 剪枝 M*N->M\n",
    "\n",
    "print(len(solution))\n",
    "printSeq(solution)\n",
    "# 用额外的一次query再问一下最终的答案\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total tokens: 1103\n",
      "total tokens: 185\n",
      "error\n"
     ]
    }
   ],
   "source": [
    "# 从M个路径里挑一个最好的来问,也可以问完之后再评估一下选最好的答案\n",
    "\n",
    "user_q = f\"\"\"There is a math_problem:\\n{question}\n",
    "\n",
    "I have broken this math problem down into a series of smaller problems and each sub-problem is solved.\n",
    "The sub-problems and their corresponding answers are as follows. (Format: Sub-problem-Id: xxx; Sub-problem: xxx; Answer: xxx.)\"\"\"\n",
    "\n",
    "for index, value in enumerate(solution[0][1]):  # 这里仅仅使用了最终得分最高的1条路径来总结得到final answer\n",
    "    user_q += f\"\"\"\\nSub-problem-Id: {index}; Sub-problem: {steps[index]}; Answer: {value}.\"\"\"\n",
    "\n",
    "Q.append({'role':'user', 'content':f\"\"\"Now that all the sub-problems have been solved, so what is the final answer?\n",
    "Please give the final answer without any additional explanation or clarification.\"\"\"})\n",
    "finalResult = askChatGPT(Q, model=GPT_MODEL, temperature=1)\n",
    "\n",
    "# 让大语言模型来判断有没有回答正确\n",
    "judgeAnswer = {'role':'user', 'content':f\"\"\"Here is a math problem with a standard answer and a student's solution. Please help me determine if the student's solution is correct.\n",
    "Problem: {question}\n",
    "\n",
    "Standard answer: {gold_answer}\n",
    "\n",
    "Answer: {finalResult}\n",
    "\n",
    "If the student's answer is correct, just output True; otherwise, just output False.\n",
    "No explanation is required.\n",
    "\"\"\"}\n",
    "\n",
    "Q_judge = [judgeAnswer]\n",
    "ifcorrect = askChatGPT(Q_judge, model=GPT_MODEL, temperature=1)  # 要么是True, 要么是False\n",
    "\n",
    "if 'True' in ifcorrect:\n",
    "    print('correct')\n",
    "    success = True  # 任务未受中断,完整地结束了,所以标记为成功\n",
    "elif 'False' in ifcorrect:\n",
    "    print('error')\n",
    "    success = True  # 任务未受中断,完整地结束了,所以标记为成功                       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Sub-problem: Calculate the total number of ways n rooks can be placed peacefully on an n x n chess board.\\nAnswer: \\\\( n^{\\\\frac{n^2}{4}} \\\\)'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "finalResult"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
